using System.IO;
using System.Threading.Tasks;
using Unity.Collections;
using Unity.Mathematics;
using UnityEngine;
using Debug = UnityEngine.Debug;

namespace GPUAnimationLib.Baker
{
    /// <summary>
    /// 网格烘焙
    /// </summary>
    public static class MeshBakerApi
    {
        /// <summary>
        /// 烘焙骨骼权重到网格中（单个顶点最大支持受4根骨骼影响）
        /// uv1 : [骨骼1索引值, 骨骼1权重值, 骨骼2索引值, 骨骼2权重值]
        /// uv2 : [骨骼3索引值, 骨骼3权重值, 骨骼4索引值, 骨骼4权重值]
        /// </summary>
        /// <param name="sourceMesh">网格</param>
        /// <param name="maxNumberOfBonesPerVertex">单个顶点最大支持受多少根骨骼影响</param>
        /// <returns>实例化的网格对象</returns>
        public static Mesh BakeBoneWeightsIntoMesh(Mesh sourceMesh, int maxNumberOfBonesPerVertex)
        {
            Mesh mesh = GameObject.Instantiate(sourceMesh);
            int vertexCount = mesh.vertices.Length;
            NativeArray<Vector4> uvs1 = new NativeArray<Vector4>(vertexCount, Allocator.Temp);
            NativeArray<Vector4> uvs2 = new NativeArray<Vector4>(vertexCount, Allocator.Temp);
            NativeArray<byte> bonesPerVertex = mesh.GetBonesPerVertex();
            NativeArray<BoneWeight1> boneWeightsSource = mesh.GetAllBoneWeights();

            int startBoneWeightIndex = 0;

            for (int vertIndex = 0; vertIndex < vertexCount; vertIndex++)
            {
                float totalWeight = 0f;
                float totalWeightCapped = 0f;
                int numberOfBonesForThisVertex = bonesPerVertex[vertIndex];

                int boneWeightIndexTemp = startBoneWeightIndex;
                for (int i = 0; i < numberOfBonesForThisVertex; i++)
                {
                    BoneWeight1 currentBoneWeight = boneWeightsSource[boneWeightIndexTemp];
                    totalWeight += currentBoneWeight.weight;
                    if (i < maxNumberOfBonesPerVertex) totalWeightCapped += currentBoneWeight.weight;
                    if (i > 0) Debug.Assert(boneWeightsSource[boneWeightIndexTemp - 1].weight >= currentBoneWeight.weight);
                    boneWeightIndexTemp++;
                }
                Debug.Assert(Mathf.Approximately(1f, totalWeight));

                float weightMultiplier = totalWeight / totalWeightCapped;
                int nbrOfBonesToBake = math.min(maxNumberOfBonesPerVertex, numberOfBonesForThisVertex);

                totalWeight = 0f;
                boneWeightIndexTemp = startBoneWeightIndex;
                float4 uv1 = float4.zero;
                float4 uv2 = float4.zero;
                for (int i = 0; i < nbrOfBonesToBake; i++)
                {
                    BoneWeight1 currentBoneWeight = boneWeightsSource[boneWeightIndexTemp];
                    float adjustedWeight = currentBoneWeight.weight * weightMultiplier; 
                    totalWeight += adjustedWeight;
                    boneWeightIndexTemp++;
                    //uv1 : [骨骼1索引值, 骨骼1权重值, 骨骼2索引值, 骨骼2权重值]
                    //uv2 : [骨骼3索引值, 骨骼3权重值, 骨骼4索引值, 骨骼4权重值]
                    if      (i == 0) uv1 = new float4(currentBoneWeight.boneIndex, adjustedWeight, uv1.z, uv1.w);
                    else if (i == 1) uv1 = new float4(uv1.x, uv1.y, currentBoneWeight.boneIndex, adjustedWeight);
                    else if (i == 2) uv2 = new float4(currentBoneWeight.boneIndex, adjustedWeight, uv2.z, uv2.w);
                    else if (i == 3) uv2 = new float4(uv2.x, uv2.y, currentBoneWeight.boneIndex, adjustedWeight);
                }
                Debug.Assert(Mathf.Approximately(1f, totalWeight));
                uvs1[vertIndex] = uv1;
                uvs2[vertIndex] = uv2;
                startBoneWeightIndex += numberOfBonesForThisVertex;
            }
            
            mesh.SetUVs(1, uvs1);
            mesh.SetUVs(2, uvs2);
            
            uvs1.Dispose();
            uvs2.Dispose();
            bonesPerVertex.Dispose();
            boneWeightsSource.Dispose();
            return mesh;
        }

        /// <summary>
        /// 获取单个顶点受影响的最大骨骼数
        /// </summary>
        public static int GetMaxBonesPerVertex(Mesh mesh)
        {
            NativeArray<byte> bonesPerVertex = mesh.GetBonesPerVertex();
            int maxCount = int.MinValue;
            foreach (var cnt in bonesPerVertex)
            {
                if (cnt > maxCount)
                    maxCount = cnt;
            }
            return maxCount;
        }
        
        /// <summary>
        /// 设置Mesh是否可读写
        /// </summary>
        /// <param name="isReadable">是否可读写</param>
        /// <param name="meshPath">Mesh文件路径</param>
        public static void SetMeshReadable(bool isReadable, string meshPath)
        {
            Task.Run(() =>
            {
                var content = File.ReadAllText(meshPath);
                string readableStr1 = $"m_IsReadable: {(isReadable ? 1 : 0)}";
                string readableStr2 = $"m_IsReadable: {(isReadable ? 0 : 1)}";
                if (content.Contains(readableStr2))
                {
                    content = content.Replace(readableStr2, readableStr1);
                    File.WriteAllText(meshPath, content);
                }
            });
        }
    }
}